package top.hjcwzx.wrench.limiter.strategy.algorithm.impl;

import top.hjcwzx.wrench.limiter.memory.ExpirableKeyStore;
import top.hjcwzx.wrench.limiter.strategy.algorithm.RateLimiterAlgorithmStrategy;

import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;

public class LeakyBucketRateLimiterStrategy extends ExpirableKeyStore<String, LeakyBucketRateLimiterStrategy.LeakyBucket>
        implements RateLimiterAlgorithmStrategy {

    private final int capacity;
    // 每秒流出多少请求
    private final double leakRate;

    public LeakyBucketRateLimiterStrategy(
            EvictionPolicy evictionPolicy,
            long expireAfterAccessMs,
            int maxCacheSize,
            int capacity,
            double leakRate) {
        super(evictionPolicy, expireAfterAccessMs, maxCacheSize);
        this.capacity = capacity;
        this.leakRate = leakRate;
    }

    @Override
    public boolean tryAcquire(String key, int permits) {
        LeakyBucket bucket = get(key);

        if (bucket == null) {
            bucket = new LeakyBucket(capacity, leakRate);
            put(key, bucket);
        }

        return bucket.tryConsume(permits);
    }

    static class LeakyBucket {
        private final int capacity;
        private final double leakRatePerSecond; // 每秒可流出数量
        private double water;                   // 当前水位
        private long lastLeakTimeNanos;

        public LeakyBucket(int capacity, double leakRatePerSecond) {
            this.capacity = capacity;
            this.leakRatePerSecond = leakRatePerSecond;
            this.water = 0;
            this.lastLeakTimeNanos = System.nanoTime();
        }

        public synchronized boolean tryConsume(int permits) {
            leak();

            if (water + permits <= capacity) {
                water += permits;
                return true;
            } else {
                return false;
            }
        }

        private void leak() {
            long now = System.nanoTime();
            double secondsElapsed = (now - lastLeakTimeNanos) / 1_000_000_000.0;
            double leaked = secondsElapsed * leakRatePerSecond;

            if (leaked > 0) {
                water = Math.max(0, water - leaked);
                lastLeakTimeNanos = now;
            }
        }
    }
}
